import torch
import models_seggpt
import tkinter as tk
from tkinter import filedialog
import cv2

SEG_MODEL = None
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def get_seggpt_model():
    global SEG_MODEL
    if SEG_MODEL is None:
        arch='seggpt_vit_large_patch16_input896x448'
        model = getattr(models_seggpt, arch)()
        model.seg_type = 'instance'
        checkpoint = torch.load('../../deep-models/seggpt_vit_large.pth', map_location=DEVICE)
        model.load_state_dict(checkpoint['model'], strict=False)
        model.to(DEVICE)
        model.eval()
        SEG_MODEL = model
    return SEG_MODEL

def get_a_video():
    root = tk.Tk()
    root.withdraw()
    video_path = filedialog.askopenfilename(
        title="选择视频文件",
        filetypes=[
            ("视频文件", "*.mp4 *.avi *.mov *.mkv"),
            ("所有文件", "*.*")
        ]
    )
    if not video_path:
        print("未选择文件")
        return (None, None, None)
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("无法打开视频文件")
        return (None, None, None)
    else:
        ret, frame0 = cap.read()
        if ret:
            print(f"成功打开视频，分辨率: {frame0.shape[1]}x{frame0.shape[0]}")
            return (cap, frame0, video_path)
        else:
            print("无法读取视频帧")
            return (None, None, None)
 